import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms, datasets, models
from torch.utils.data import DataLoader
import torch.optim as optim
from tqdm import tqdm
import matplotlib.pyplot as plt
import os

# 定义数据预处理和加载器
input_size = 160
scale_factor = 4

transform_train = transforms.Compose([
    transforms.Resize((input_size, input_size)),
    transforms.Grayscale(num_output_channels=3),
    transforms.ToTensor(),
    transforms.Lambda(lambda img_tensor: torch.round(img_tensor * 255 / scale_factor) / 1.0),
    transforms.Normalize(mean=[32, 32, 32], std=[1, 1, 1]),
])

transform_val = transforms.Compose([
    transforms.Resize((input_size, input_size)),
    transforms.Grayscale(num_output_channels=3),
    transforms.ToTensor(),
    transforms.Lambda(lambda img_tensor: torch.round(img_tensor * 255 / scale_factor) / 1.0),
    transforms.Normalize(mean=[32, 32, 32], std=[1, 1, 1]),
])

batch_size = 32

class QuantizedReLU(nn.Module):
    def __init__(self, scale_factor):
        super(QuantizedReLU, self).__init__()
        self.scale_factor = scale_factor

    def forward(self, x):
        x = x / self.scale_factor
        x = F.relu(x)
        return x

scale_factor_weight = 31.0
class QuantizedResNet(nn.Module):
    def __init__(self, scale_factor_weight):
        super(QuantizedResNet, self).__init__()
        self.scale_factor = scale_factor_weight
        self.original_model = models.resnet18(pretrained=True)
        self.feature_maps = []  # 用于存储每层特征图的列表

        # 创建一个新的模型，并复制原始模型的各个部分
        self.modified_model = self.copy_model(self.original_model)

        # 替换全连接层
        num_ftrs = self.original_model.fc.in_features
        self.fc = nn.Linear(num_ftrs, 196)  # 修改输出层为196类

    def copy_model(self, original_model):
        # 创建一个新的模型
        new_model = nn.Sequential()

        for name, module in original_model.named_children():
            if isinstance(module, nn.ReLU):
                # 将ReLU替换为QuantizedReLU
                new_model.add_module(name, QuantizedReLU(self.scale_factor))
            else:
                # 复制原始模型的其他部分
                new_model.add_module(name, module)

        return new_model

    def forward(self, x):
        self.feature_maps = []  # 清空特征图列表
        x = x * self.scale_factor  # 放缩输入到网络
        for name, module in self.modified_model.named_children():
            x = module(x)
            if isinstance(module, QuantizedReLU):
                self.feature_maps.append(x)  # 记录QuantizedReLU层的输出作为特征图

        # 添加一个展平层，以适应全连接层的输入要求
        x = torch.flatten(x, 1)  # 将张量展平，保持第一维不变

        x = self.fc(x)  # 应用全连接层

        return x


train_set = datasets.ImageFolder(root='../data/sdfcar/train', transform=transform_train)
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=0)

val_set = datasets.ImageFolder(root='../data/sdfcar/valid', transform=transform_val)
val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=0)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = QuantizedResNet(scale_factor)
num_ftrs = model.modified_model.fc.in_features
model.modified_model.fc = nn.Linear(num_ftrs, 196)  # 修改输出层为196类
model = model.to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=0.0002)

# 训练参数
num_epochs = 2
train_accuracy_history = []
val_accuracy_history = []

# 创建保存特征图的目录
feature_maps_dir = 'feature_maps'
os.makedirs(feature_maps_dir, exist_ok=True)

# 训练循环
for epoch in range(num_epochs):
    print(f'Epoch {epoch + 1}/{num_epochs}')
    model.train()

    running_loss = 0.0
    correct_train = 0
    total_train = 0

    for inputs, labels in tqdm(train_loader):
        inputs = inputs.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

        _, predicted = torch.max(outputs, 1)
        total_train += labels.size(0)
        correct_train += (predicted == labels).sum().item()

    epoch_train_accuracy = correct_train / total_train
    train_accuracy_history.append(epoch_train_accuracy)

    print(f'Training Loss: {running_loss / len(train_loader):.4f}, Training Accuracy: {epoch_train_accuracy * 100:.2f}%')

    # 验证模型
    model.eval()
    correct_val = 0
    total_val = 0

    with torch.no_grad():
        for inputs, labels in tqdm(val_loader):
            inputs = inputs.to(device)
            labels = labels.to(device)

            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            total_val += labels.size(0)
            correct_val += (predicted == labels).sum().item()

    epoch_val_accuracy = correct_val / total_val
    val_accuracy_history.append(epoch_val_accuracy)

    print(f'Validation Accuracy: {epoch_val_accuracy * 100:.2f}%')

    # 保存每层特征图
    for idx, feature_map in enumerate(model.feature_maps, 1):
        feature_map_path = os.path.join(feature_maps_dir, f'epoch_{epoch + 1}_layer_{idx}.pt')
        torch.save(feature_map, feature_map_path)

    # 保存模型
    torch.save(model.state_dict(), f'models/model_epoch_{epoch + 1}.pth')

# 保存训练准确率和验证准确率的曲线图
plt.figure(figsize=(10, 5))
plt.plot(range(1, num_epochs + 1), train_accuracy_history, label='Training Accuracy')
plt.plot(range(1, num_epochs + 1), val_accuracy_history, label='Validation Accuracy')
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.title('Training and Validation Accuracy over Epochs')
plt.legend()
plt.grid(True)
plt.savefig('accuracy_plot.png')
plt.show()
